Cutlery/calibre
1
0
Fork 0
mirror of https://github.com/kovidgoyal/calibre.git synced 2025-07-09 03:04:10 -04:00
Code Issues Packages Projects Releases Wiki Activity

DRYer

Browse Source
This commit is contained in:
Kovid Goyal 2016-01-18 08:10:12 +05:30
parent a239bf686a
commit 4f5146e2ea
1 changed files with 100 additions and 125 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

117
src/calibre/devices/winusb.py
View File

@ -6,7 +6,6 @@ from __future__ import (unicode_literals, division, absolute_import,
print_function) print_function)
import os, string, _winreg as winreg, re, time, sys import os, string, _winreg as winreg, re, time, sys
from contextlib import contextmanager
from ctypes import ( from ctypes import (
Structure, POINTER, c_ubyte, windll, byref, c_void_p, WINFUNCTYPE, Structure, POINTER, c_ubyte, windll, byref, c_void_p, WINFUNCTYPE,
WinError, get_last_error, sizeof, c_wchar, create_string_buffer, cast, WinError, get_last_error, sizeof, c_wchar, create_string_buffer, cast,
@ -292,11 +291,42 @@ CM_Get_Device_Interface_List = cwrap('CM_Get_Device_Interface_ListW', CONFIGRET,
# }}} # }}}
# Utility functions {{{ # Utility functions {{{
@contextmanager class DeviceSet(object):
def get_device_set(guid=byref(GUID_DEVINTERFACE_VOLUME), enumerator=None, flags=DIGCF_PRESENT | DIGCF_DEVICEINTERFACE):
dev_list = SetupDiGetClassDevs(guid, enumerator, None, flags) def __init__(self, guid=GUID_DEVINTERFACE_VOLUME, enumerator=None, flags=DIGCF_PRESENT | DIGCF_DEVICEINTERFACE):
yield dev_list self.guid_ref, self.enumerator, self.flags = (None if guid is None else byref(guid)), enumerator, flags
SetupDiDestroyDeviceInfoList(dev_list) self.dev_list = SetupDiGetClassDevs(self.guid_ref, self.enumerator, None, self.flags)
def __del__(self):
SetupDiDestroyDeviceInfoList(self.dev_list)
del self.dev_list
def interfaces(self, ignore_errors=False):
interface_data = SP_DEVICE_INTERFACE_DATA()
interface_data.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA)
buf = None
i = -1
while True:
i += 1
if not SetupDiEnumDeviceInterfaces(self.dev_list, None, self.guid_ref, i, byref(interface_data)):
break
try:
buf, devinfo, devpath = get_device_interface_detail_data(self.dev_list, byref(interface_data), buf)
except WindowsError:
if ignore_errors:
continue
raise
yield devinfo, devpath
def devices(self):
devinfo = SP_DEVINFO_DATA()
devinfo.cbSize = sizeof(SP_DEVINFO_DATA)
i = -1
while True:
i += 1
if not SetupDiEnumDeviceInfo(self.dev_list, i, byref(devinfo)):
break
yield self.dev_list, devinfo
def iterchildren(parent_devinst): def iterchildren(parent_devinst):
@ -450,18 +480,8 @@ def get_removable_drives(debug=False): # {{{
if not drive_map: if not drive_map:
raise NoRemovableDrives('No removable drives found!') raise NoRemovableDrives('No removable drives found!')
buf = None ans, buf = {}, None
pbuf = create_unicode_buffer(512) for devinfo, devpath in DeviceSet().interfaces():
ans = {}
with get_device_set() as dev_list:
interface_data = SP_DEVICE_INTERFACE_DATA()
interface_data.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA)
i = -1
while True:
i += 1
if not SetupDiEnumDeviceInterfaces(dev_list, None, byref(GUID_DEVINTERFACE_VOLUME), i, byref(interface_data)):
break
buf, devinfo, devpath = get_device_interface_detail_data(dev_list, byref(interface_data), buf)
candidates = [] candidates = []
# Get the devpaths for all parents of this device. This is not # Get the devpaths for all parents of this device. This is not
# actually necessary on Vista+, so we simply ignore any windows API # actually necessary on Vista+, so we simply ignore any windows API
@ -473,7 +493,7 @@ def get_removable_drives(debug=False): # {{{
except WindowsError: except WindowsError:
break break
try: try:
devid, pbuf = get_device_id(parent, buf=pbuf) devid, buf = get_device_id(parent, buf=buf)
except WindowsError: except WindowsError:
break break
candidates.append(devid) candidates.append(devid)
@ -497,15 +517,7 @@ def get_drive_letters_for_device(vendor_id, product_id, bcd=None, debug=False):
ans = [] ans = []
# First search for a device matching the specified USB ids # First search for a device matching the specified USB ids
with get_device_set(enumerator='USB', flags=DIGCF_PRESENT | DIGCF_ALLCLASSES) as dev_list: for dev_list, devinfo in DeviceSet(enumerator='USB', flags=DIGCF_PRESENT | DIGCF_ALLCLASSES).devices():
devinfo = SP_DEVINFO_DATA()
devinfo.cbSize = sizeof(SP_DEVINFO_DATA)
i = -1
found_at = None
while True:
i += 1
if not SetupDiEnumDeviceInfo(dev_list, i, byref(devinfo)):
break
rbuf, devid = get_device_registry_property(dev_list, byref(devinfo), buf=rbuf) rbuf, devid = get_device_registry_property(dev_list, byref(devinfo), buf=rbuf)
if devid: if devid:
m = devid_pat().search(devid[0]) m = devid_pat().search(devid[0])
@ -516,9 +528,8 @@ def get_drive_letters_for_device(vendor_id, product_id, bcd=None, debug=False):
except Exception: except Exception:
continue continue
if vid == vendor_id and pid == product_id and (bcd is None or (bcd and rev in bcd)): if vid == vendor_id and pid == product_id and (bcd is None or (bcd and rev in bcd)):
found_at = i - 1
break break
if found_at is None: else:
if debug: if debug:
prints('Could not find device matching vid=0x%x pid=0x%x' % (vendor_id, product_id)) prints('Could not find device matching vid=0x%x pid=0x%x' % (vendor_id, product_id))
return ans return ans
@ -577,16 +588,7 @@ def get_storage_number_map(drive_types=(DRIVE_REMOVABLE, DRIVE_FIXED), debug=Fal
def find_drive(devinst, storage_number_map, debug=False): def find_drive(devinst, storage_number_map, debug=False):
buf = None for devinfo, devpath in DeviceSet(GUID_DEVINTERFACE_DISK).interfaces():
interface_data = SP_DEVICE_INTERFACE_DATA()
interface_data.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA)
with get_device_set(guid=byref(GUID_DEVINTERFACE_DISK)) as dev_list:
i = -1
while True:
i += 1
if not SetupDiEnumDeviceInterfaces(dev_list, None, byref(GUID_DEVINTERFACE_DISK), i, byref(interface_data)):
break
buf, devinfo, devpath = get_device_interface_detail_data(dev_list, byref(interface_data), buf=buf)
if devinfo.DevInst == devinst: if devinfo.DevInst == devinst:
storage_number = get_storage_number(devpath) storage_number = get_storage_number(devpath)
drive_letter = storage_number_map.get(storage_number) drive_letter = storage_number_map.get(storage_number)
@ -596,16 +598,8 @@ def find_drive(devinst, storage_number_map, debug=False):
# }}} # }}}
def get_usb_devices(): # {{{ def get_usb_devices(): # {{{
buf = None ans, buf = [], None
ans = [] for dev_list, devinfo in DeviceSet(guid=None, enumerator='USB', flags=DIGCF_PRESENT | DIGCF_ALLCLASSES).devices():
with get_device_set(guid=None, enumerator='USB', flags=DIGCF_PRESENT | DIGCF_ALLCLASSES) as dev_list:
devinfo = SP_DEVINFO_DATA()
devinfo.cbSize = sizeof(SP_DEVINFO_DATA)
i = -1
while True:
i += 1
if not SetupDiEnumDeviceInfo(dev_list, i, byref(devinfo)):
break
buf, devid = get_device_registry_property(dev_list, byref(devinfo), buf=buf) buf, devid = get_device_registry_property(dev_list, byref(devinfo), buf=buf)
if devid: if devid:
ans.append(devid[0].lower()) ans.append(devid[0].lower())
@ -614,14 +608,7 @@ def get_usb_devices(): # {{{
def is_usb_device_connected(vendor_id, product_id): # {{{ def is_usb_device_connected(vendor_id, product_id): # {{{
buf = None buf = None
with get_device_set(guid=None, enumerator='USB', flags=DIGCF_PRESENT | DIGCF_ALLCLASSES) as dev_list: for dev_list, devinfo in DeviceSet(guid=None, enumerator='USB', flags=DIGCF_PRESENT | DIGCF_ALLCLASSES).devices():
devinfo = SP_DEVINFO_DATA()
devinfo.cbSize = sizeof(SP_DEVINFO_DATA)
i = -1
while True:
i += 1
if not SetupDiEnumDeviceInfo(dev_list, i, byref(devinfo)):
break
buf, devid = get_device_registry_property(dev_list, byref(devinfo), buf=buf) buf, devid = get_device_registry_property(dev_list, byref(devinfo), buf=buf)
if devid: if devid:
m = devid_pat().search(devid[0]) m = devid_pat().search(devid[0])
@ -663,7 +650,6 @@ def eject_drive(drive_letter): # {{{
def devinst_from_device_number(drive_letter, device_number): def devinst_from_device_number(drive_letter, device_number):
drive_root = drive_letter + ':' + os.sep drive_root = drive_letter + ':' + os.sep
buf = create_unicode_buffer(512) buf = create_unicode_buffer(512)
bbuf = None
sdn = STORAGE_DEVICE_NUMBER() sdn = STORAGE_DEVICE_NUMBER()
drive_type = GetDriveType(drive_root) drive_type = GetDriveType(drive_root)
QueryDosDevice(drive_letter + ':', buf, len(buf)) QueryDosDevice(drive_letter + ':', buf, len(buf))
@ -676,18 +662,7 @@ def devinst_from_device_number(drive_letter, device_number):
guid = GUID_DEVINTERFACE_CDROM guid = GUID_DEVINTERFACE_CDROM
else: else:
raise ValueError('Unknown drive_type: %d' % drive_type) raise ValueError('Unknown drive_type: %d' % drive_type)
with get_device_set(guid=guid) as dev_list: for devinfo, devpath in DeviceSet(guid=guid).interfaces(ignore_errors=True):
interface_data = SP_DEVICE_INTERFACE_DATA()
interface_data.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA)
i = -1
while True:
i += 1
if not SetupDiEnumDeviceInterfaces(dev_list, None, byref(guid), i, byref(interface_data)):
break
try:
bbuf, devinfo, devpath = get_device_interface_detail_data(dev_list, byref(interface_data), bbuf)
except WindowsError:
continue
handle = CreateFile(devpath, 0, FILE_SHARE_READ | FILE_SHARE_WRITE, None, OPEN_EXISTING, 0, None) handle = CreateFile(devpath, 0, FILE_SHARE_READ | FILE_SHARE_WRITE, None, OPEN_EXISTING, 0, None)
try: try:
DeviceIoControl(handle, IOCTL_STORAGE_GET_DEVICE_NUMBER, None, 0, byref(sdn), sizeof(STORAGE_DEVICE_NUMBER), None, None) DeviceIoControl(handle, IOCTL_STORAGE_GET_DEVICE_NUMBER, None, 0, byref(sdn), sizeof(STORAGE_DEVICE_NUMBER), None, None)