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

225
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,41 +480,31 @@ 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 = {} candidates = []
with get_device_set() as dev_list: # Get the devpaths for all parents of this device. This is not
interface_data = SP_DEVICE_INTERFACE_DATA() # actually necessary on Vista+, so we simply ignore any windows API
interface_data.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA) # failures.
i = -1 parent = DEVINST(devinfo.DevInst)
while True: while True:
i += 1 try:
if not SetupDiEnumDeviceInterfaces(dev_list, None, byref(GUID_DEVINTERFACE_VOLUME), i, byref(interface_data)): CM_Get_Parent(byref(parent), parent, 0)
except WindowsError:
break break
buf, devinfo, devpath = get_device_interface_detail_data(dev_list, byref(interface_data), buf) try:
candidates = [] devid, buf = get_device_id(parent, buf=buf)
# Get the devpaths for all parents of this device. This is not except WindowsError:
# actually necessary on Vista+, so we simply ignore any windows API break
# failures. candidates.append(devid)
parent = DEVINST(devinfo.DevInst) candidates.append(devpath)
while True:
try:
CM_Get_Parent(byref(parent), parent, 0)
except WindowsError:
break
try:
devid, pbuf = get_device_id(parent, buf=pbuf)
except WindowsError:
break
candidates.append(devid)
candidates.append(devpath)
drive_letter = drive_letter_from_volume_devpath(devpath, drive_map) drive_letter = drive_letter_from_volume_devpath(devpath, drive_map)
if drive_letter: if drive_letter:
ans[drive_letter] = candidates ans[drive_letter] = candidates
if debug: if debug:
prints('Found volume with device path:', devpath, ' Drive letter:', drive_letter, 'Is removable:', drive_letter in ans) prints('Found volume with device path:', devpath, ' Drive letter:', drive_letter, 'Is removable:', drive_letter in ans)
return ans return ans
# }}} # }}}
def get_drive_letters_for_device(vendor_id, product_id, bcd=None, debug=False): # {{{ def get_drive_letters_for_device(vendor_id, product_id, bcd=None, debug=False): # {{{
@ -497,28 +517,19 @@ 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() rbuf, devid = get_device_registry_property(dev_list, byref(devinfo), buf=rbuf)
devinfo.cbSize = sizeof(SP_DEVINFO_DATA) if devid:
i = -1 m = devid_pat().search(devid[0])
found_at = None if m is None:
while True: continue
i += 1 try:
if not SetupDiEnumDeviceInfo(dev_list, i, byref(devinfo)): vid, pid, rev = map(lambda x:int(x, 16), m.group(1, 2, 3))
except Exception:
continue
if vid == vendor_id and pid == product_id and (bcd is None or (bcd and rev in bcd)):
break break
rbuf, devid = get_device_registry_property(dev_list, byref(devinfo), buf=rbuf) else:
if devid:
m = devid_pat().search(devid[0])
if m is None:
continue
try:
vid, pid, rev = map(lambda x:int(x, 16), m.group(1, 2, 3))
except Exception:
continue
if vid == vendor_id and pid == product_id and (bcd is None or (bcd and rev in bcd)):
found_at = i - 1
break
if found_at is None:
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,61 +588,37 @@ 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() if devinfo.DevInst == devinst:
interface_data.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA) storage_number = get_storage_number(devpath)
with get_device_set(guid=byref(GUID_DEVINTERFACE_DISK)) as dev_list: drive_letter = storage_number_map.get(storage_number)
i = -1 if drive_letter:
while True: return drive_letter
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:
storage_number = get_storage_number(devpath)
drive_letter = storage_number_map.get(storage_number)
if drive_letter:
return drive_letter
# }}} # }}}
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: buf, devid = get_device_registry_property(dev_list, byref(devinfo), buf=buf)
devinfo = SP_DEVINFO_DATA() if devid:
devinfo.cbSize = sizeof(SP_DEVINFO_DATA) ans.append(devid[0].lower())
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)
if devid:
ans.append(devid[0].lower())
return ans return ans
# }}} # }}}
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() buf, devid = get_device_registry_property(dev_list, byref(devinfo), buf=buf)
devinfo.cbSize = sizeof(SP_DEVINFO_DATA) if devid:
i = -1 m = devid_pat().search(devid[0])
while True: if m is not None:
i += 1 try:
if not SetupDiEnumDeviceInfo(dev_list, i, byref(devinfo)): vid, pid = map(lambda x: int(x, 16), m.group(1, 2))
break except Exception:
buf, devid = get_device_registry_property(dev_list, byref(devinfo), buf=buf) continue
if devid: if vid == vendor_id and pid == product_id:
m = devid_pat().search(devid[0]) return True
if m is not None:
try:
vid, pid = map(lambda x: int(x, 16), m.group(1, 2))
except Exception:
continue
if vid == vendor_id and pid == product_id:
return True
return False return False
# }}} # }}}
@ -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,25 +662,14 @@ 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() handle = CreateFile(devpath, 0, FILE_SHARE_READ | FILE_SHARE_WRITE, None, OPEN_EXISTING, 0, None)
interface_data.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA) try:
i = -1 DeviceIoControl(handle, IOCTL_STORAGE_GET_DEVICE_NUMBER, None, 0, byref(sdn), sizeof(STORAGE_DEVICE_NUMBER), None, None)
while True: finally:
i += 1 CloseHandle(handle)
if not SetupDiEnumDeviceInterfaces(dev_list, None, byref(guid), i, byref(interface_data)): if sdn.DeviceNumber == device_number:
break return devinfo.DevInst
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)
try:
DeviceIoControl(handle, IOCTL_STORAGE_GET_DEVICE_NUMBER, None, 0, byref(sdn), sizeof(STORAGE_DEVICE_NUMBER), None, None)
finally:
CloseHandle(handle)
if sdn.DeviceNumber == device_number:
return devinfo.DevInst
# }}} # }}}